from datetime import datetime
import matplotlib.pyplot as plt
import akshare as ak
import pandas as pd
from datetime import datetime, timedelta
import numpy as np  
import talib  

plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

# 获取股票后复权数据
df = ak.stock_zh_a_hist(symbol="000001", period="daily", 
                        start_date="20200101", end_date='20200315', adjust="qfq").iloc[:, :6]
# 处理字段命名
df.columns = ['date', 'open', 'close', 'high', 'low', 'volume']
# 将date设为索引
df.index = pd.to_datetime(df['date'])


# 计算短期（例如5天）和长期（例如20天）SMA  
short_window = 5  
long_window = 20  
df['SMA_short'] = talib.SMA(df['close'], timeperiod=short_window)  
df['SMA_long'] = talib.SMA(df['close'], timeperiod=long_window)  
  

 
# 初始化标记列表  
cross_up_positions = []  # 金叉位置  
cross_down_positions = []  # 死叉位置  
  
# 遍历数据以找到交叉点  
for i in range(1, len(df)):  
    if df['SMA_short'].iloc[i] > df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] <= df['SMA_long'].iloc[i-1]:  
        cross_up_positions.append((df.index[i], df['SMA_short'].iloc[i]))  # 添加交叉点和对应的值  
    if df['SMA_short'].iloc[i] < df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] >= df['SMA_long'].iloc[i-1]:  
        cross_down_positions.append((df.index[i], df['SMA_short'].iloc[i]))  # 添加交叉点和对应的值  
  
# 绘制股价和均线  
plt.figure(figsize=(12, 6))  
plt.plot(df.index, df['close'], label='Close Price', alpha=0.7)  
plt.plot(df.index, df['SMA_short'], label=f'{short_window}-day SMA', color='red', alpha=0.9)  
plt.plot(df.index, df['SMA_long'], label=f'{long_window}-day SMA', color='blue', alpha=0.9)  
"""   
# 绘制交叉点上的箭头  
for pos, value in cross_up_positions:  
    # 绘制向上的小箭头  
    plt.annotate('', xy=(pos, value), xytext=(10, 0), # xytext是箭头尖端的位置  
                 textcoords="offset points",  
                 arrowprops=dict(arrowstyle="->", color='green'))  
  
for pos, value in cross_down_positions:  
    # 绘制向下的小箭头  
    plt.annotate('', xy=(pos, value), xytext=(0, -10), # xytext是箭头尖端的位置  
                 textcoords="offset points",  
                 arrowprops=dict(arrowstyle="->", color='red'))  
"""

# 绘制交叉点上的实心小箭头（用三角形表示）  
def draw_solid_arrow(ax, x, y, direction, color, size=5):  
    # direction应该是'up'或'down'  
    dx, dy = 0, size if direction == 'up' else -size  
    if direction == 'up':  
        # 向上小三角形  
        ax.scatter([x], [y], marker=(5, 1, dy), color=color, s=size**2)  
    else:  
        # 向下小三角形  
        ax.scatter([x], [y], marker=(5, 2, -dy), color=color, s=size**2)  
  
for pos, value in cross_up_positions:  
    draw_solid_arrow(plt.gca(), pos, value, 'up', 'green', size=8)  
  
for pos, value in cross_down_positions:  
    draw_solid_arrow(plt.gca(), pos, value, 'down', 'red', size=8)  


# 设置图例和标签  
plt.legend()  
plt.title('Stock Price with SMA Crossover Signals')  
plt.xlabel('Date')  
plt.ylabel('Price')  
plt.grid(True)  
  
# 显示图表  
plt.show()

""" # 初始化标记列表  
cross_up = []  # 金叉（短期线上穿长期线）  
cross_down = []  # 死叉（短期线下穿长期线）  
  
# 遍历数据以找到交叉点  
for i in range(1, len(df)):  
    # 注意：使用.iloc[]来按位置访问数据  
    if df['SMA_short'].iloc[i] > df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] <= df['SMA_long'].iloc[i-1]:  
        cross_up.append(df.index[i])  
    if df['SMA_short'].iloc[i] < df['SMA_long'].iloc[i] and df['SMA_short'].iloc[i-1] >= df['SMA_long'].iloc[i-1]:  
        cross_down.append(df.index[i])
  
# 绘制股价和均线  
plt.figure(figsize=(12, 6))  
plt.plot(df.index, df['close'], label='Close Price', alpha=0.7)  
plt.plot(df.index, df['SMA_short'], label=f'{short_window}-day SMA', color='red', alpha=0.9)  
plt.plot(df.index, df['SMA_long'], label=f'{long_window}-day SMA', color='blue', alpha=0.9)  
  
# 在交叉点添加标记  
for date in cross_up:  
    plt.axvline(x=date, color='green', linestyle='--', linewidth=1)  
for date in cross_down:  
    plt.axvline(x=date, color='red', linestyle='--', linewidth=1)  
  
# 设置图例和标签  
plt.legend()  
plt.title('Stock Price with SMA Crossover Signals')  
plt.xlabel('Date')  
plt.ylabel('Price')  
plt.grid(True)  
  
# 显示图表  
plt.show()
 """